/*
 * Copyright (c) HiSilicon (Shanghai) Technologies Co., Ltd. 2023-2023. All rights reserved.
 */

#ifndef BURNCTRLDFU_H
#define BURNCTRLDFU_H

#include <QObject>
#include <windows.h>

#include "BurnCtrl.h"
#include "libusb.h"
#include "BurnToolCommon/usbDfu.h"

#define QUIRK_POLLTIMEOUT  (1<<0)
#define QUIRK_FORCE_DFU11  (1<<1)
#define QUIRK_UTF8_SERIAL  (1<<2)
#define QUIRK_DFUSE_LAYOUT (1<<3)
#define QUIRK_DFUSE_LEAVE  (1<<4)

namespace BurnCtrlDfuClass {
    const char DFUSE_READABLE = 1;
    const char DFUSE_ERASABLE = 2;
    const char DFUSE_WRITEABLE = 4;

    const uint16_t VENDOR_OPENMOKO = 0x1d50;    /* Openmoko Freerunner / GTA02 */
    const uint16_t VENDOR_FIC = 0x1457;         /* Openmoko Freerunner / GTA02 */
    const uint16_t VENDOR_VOTI = 0x16c0;        /* OpenPCD Reader */
    const uint16_t VENDOR_LEAFLABS = 0x1eaf;    /* Maple */
    const uint16_t VENDOR_SIEMENS = 0x0908;     /* Siemens AG */
    const uint16_t VENDOR_MIDIMAN = 0x0763;     /* Midiman */
    const uint16_t VENDOR_GIGADEVICE = 0x0763;  /* GigaDevice */

    const uint16_t PRODUCT_FREERUNNER_FIRST = 0x5117;
    const uint16_t PRODUCT_FREERUNNER_LAST = 0x5126;
    const uint16_t PRODUCT_SIMTRACE = 0x0762;
    const uint16_t PRODUCT_OPENPCD = 0x076b;
    const uint16_t PRODUCT_OPENPICC = 0x076c;
    const uint16_t PRODUCT_MAPLE3 = 0x0003;     /* rev 3 and 5 */
    const uint16_t PRODUCT_PXM40 = 0x02c4;      /* Siemens AG, PXM 40 */
    const uint16_t PRODUCT_PXM50 = 0x02c5;      /* Siemens AG, PXM 50 */
    const uint16_t PRODUCT_TRANSIT = 0x2806;    /* M-Audio Transit (Midiman) */
    const uint16_t PRODUCT_GD32 = 0x0189;       /* GigaDevice GD32VF103 rev 1 */

    const unsigned int DEFAULT_POLLTIMEOUT = 5;     /* Fallback value, works for OpenMoko */

    /* DFU states */
    const unsigned char STATE_APP_IDLE = 0x00;
    const unsigned char STATE_APP_DETACH = 0x01;
    const unsigned char STATE_DFU_IDLE = 0x02;
    const unsigned char STATE_DFU_DOWNLOAD_SYNC = 0x03;
    const unsigned char STATE_DFU_DOWNLOAD_BUSY = 0x04;
    const unsigned char STATE_DFU_DOWNLOAD_IDLE = 0x05;
    const unsigned char STATE_DFU_MANIFEST_SYNC = 0x06;
    const unsigned char STATE_DFU_MANIFEST = 0x07;
    const unsigned char STATE_DFU_MANIFEST_WAIT_RESET = 0x08;
    const unsigned char STATE_DFU_UPLOAD_IDLE = 0x09;
    const unsigned char STATE_DFU_ERROR = 0x0a;

    /* DFU status */
    const unsigned char DFU_STATUS_OK_VALUE = 0x00;
    const unsigned char DFU_STATUS_ERROR_TARGET = 0x01;
    const unsigned char DFU_STATUS_ERROR_FILE = 0x02;
    const unsigned char DFU_STATUS_ERROR_WRITE = 0x03;
    const unsigned char DFU_STATUS_ERROR_ERASE = 0x04;
    const unsigned char DFU_STATUS_ERROR_CHECK_ERASED = 0x05;
    const unsigned char DFU_STATUS_ERROR_PROG = 0x06;
    const unsigned char DFU_STATUS_ERROR_VERIFY = 0x07;
    const unsigned char DFU_STATUS_ERROR_ADDRESS = 0x08;
    const unsigned char DFU_STATUS_ERROR_NOTDONE = 0x09;
    const unsigned char DFU_STATUS_ERROR_FIRMWARE = 0x0a;
    const unsigned char DFU_STATUS_ERROR_VENDOR = 0x0b;
    const unsigned char DFU_STATUS_ERROR_USBR = 0x0c;
    const unsigned char DFU_STATUS_ERROR_POR = 0x0d;
    const unsigned char DFU_STATUS_ERROR_UNKNOWN = 0x0e;
    const unsigned char DFU_STATUS_ERROR_STALLEDPKT = 0x0f;

    /* DFU commands */
    const uint8_t DFU_DETACH = 0;
    const uint8_t DFU_DNLOAD = 1;
    const uint8_t DFU_UPLOAD = 2;
    const uint8_t DFU_GETSTATUS = 3;
    const uint8_t DFU_CLRSTATUS = 4;
    const uint8_t DFU_GETSTATE = 5;
    const uint8_t DFU_ABORT = 6;

    /* DFU interface */
    const unsigned int DFU_IF_F_DFU = 0x0001;    /* DFU Mode, (not Runtime) */
    const unsigned int DFU_IF_F_ALT = 0x0002;    /* Multiple alternate settings */

    const unsigned int DFU_TIMEOUT = 5000;
    const unsigned int MAX_DESC_STR_LEN = 254;
    const unsigned int TRANSFER_SIZE = 4096;

    enum class SuffixReq {
        NO_SUFFIX,
        NEEDS_SUFFIX,
        MAYBE_SUFFIX
    };

    enum class PrefixReq {
        NO_PREFIX,
        NEEDS_PREFIX,
        MAYBE_PREFIX
    };

    enum class PrefixType {
        ZERO_PREFIX,
        LMDFU_PREFIX,
        LPCDFU_UNENCRYPTED_PREFIX
    };

    enum class DfuseCommand {
        SET_ADDRESS, ERASE_PAGE, MASS_ERASE, READ_UNPROTECT
    };

struct DfuFile {
    /* File name */
    const char *name;
    /* Pointer to file loaded into memory */
    uint8_t *firmware = nullptr;
    /* Different sizes */
    struct {
        off_t total;
        int prefix;
        int suffix;
    } size;
    /* From prefix fields */
    uint32_t lmdfuAddress;
    /* From prefix fields */
    uint32_t prefixType;

    /* From DFU suffix fields */
    uint32_t dwCRC;
    uint16_t bcdDFU;
    uint16_t idVendor;
    uint16_t idProduct;
    uint16_t bcdDevice;
    uint32_t index;
    uint32_t fileSize;
};

struct DfuIf {
    struct UsbDfu::UsbDfuFuncDescriptor funcDfu;
    uint16_t quirks;
    uint16_t busnum;
    uint16_t devnum;
    uint16_t vendor;
    uint16_t product;
    uint16_t bcdDevice;
    uint8_t configuration;
    uint8_t target;
    uint8_t altsetting;
    uint8_t flags;
    uint8_t bMaxPacketSize0;
    CHAR *altName;
    CHAR *serialName;
    libusb_device *dev;
    libusb_device_handle *devHandle;
    struct DfuIf *next;
    struct memsegment *memLayout; /* for DfuSe */
};

struct memsegment {
    unsigned int start;
    unsigned int end;
    int pagesize;
    int memtype;
    struct memsegment *next;
};

enum class DfuMode {
    MODE_NONE,
    MODE_VERSION,
    MODE_LIST,
    MODE_DETACH,
    MODE_UPLOAD,
    MODE_DOWNLOAD
};

/* This is based off of DFU_GETSTATUS
 *
 *  1 unsigned byte bStatus
 *  3 unsigned byte bwPollTimeout
 *  1 unsigned byte bState
 *  1 unsigned byte iString
*/

struct DfuStatus {
    unsigned char bStatus;
    unsigned int  bwPollTimeout;
    unsigned char bState;
    unsigned char iString;
};

class BurnCtrlDfu : public BurnCtrl {
    Q_OBJECT
public:
    explicit BurnCtrlDfu(ChipType type = ChipType::USBBS25, QObject *parent = nullptr);
    ~BurnCtrlDfu() override;
    int Connect();
    int ConnectModeHidDfu();
    int ConnectSub();
    void LoadFile(QString filePath, unsigned int index, unsigned int size);
    void DoDfuDownLoad();
    void SetMode(DfuMode mode);
    void StartConnect() override;
    void CommonAction1() override;
    void CommonAction2() override;
    void CommonAction3() override;
    void CommonAction2Sub(U32 &ret);
    void CommonAction2SubAdd();
    U32 SendBurnCmd() override;
    void WriteToPort(const char *data, int len) override;
    void InitBurnCmdByArray(const QMap<int, BIN_INFO> &infoArray)  override;
    void OnlyEraseAllSendPayload(QString &strSendData) const;
    void NotOnlyEraseAllSendPayload(QString &strSendData) const;
    U32 StartBurn() override;
    void ResetUIAndFSM() override;
    void SetEraseMode(int eraseMode) override;
    void ResetByCmd() override;
    bool ConfigComPort() override;
    bool IsBusy() override;
    bool CycleUsbDevice() const;
    void EnterDfuMode();
protected:
    int DfuStateAction();
    int DfuStateActionSub();
    int WaitDfuDetach();
    int DfuDetach(libusb_device_handle *device,
        const unsigned short interface,
        const unsigned short timeout) const;
    int DfuDownload(libusb_device_handle *device,
        const unsigned short interface,
        const unsigned short length,
        const unsigned short wValue,
        PBYTE data) const;
    int DfuGetStatus(DfuIf *&dif, struct DfuStatus &status, int timeout = 5000) const;
    int DfuClearStatus(libusb_device_handle *device,
        const unsigned short interface) const;
    int DfuGetState(libusb_device_handle *device,
        const unsigned short interface) const;
    int DfuAbort(libusb_device_handle *device,
        const unsigned short interface) const;
    const char *DfuStateToString(int state) const;
    int ProbePrefix(DfuFile &mFile) const;
    const char *DfuStatusToString(int status) const;
    int DfuLoadFile(DfuFile &mFile, SuffixReq checkSuffix, PrefixReq checkPrefix,
                    unsigned int index = 0, unsigned int fileSize = 0);
    int PartOfDfuLoadFile(DfuFile &mFile, unsigned int index, unsigned int fileSize);
    int DfuLoadFileSub(DfuFile &mFile, const SuffixReq &checkSuffix, const PrefixReq &checkPrefix);

    void DfuProgressBar(const char *desc, unsigned long long curr, unsigned long long max);
    int DfuloadDoDnload(DfuIf *&dif, int xferSize, DfuFile &mFile);
    int FindDescriptor(const uint8_t *descList, int listLen, uint8_t descType, PVOID resBuf, int resSize);
    int GetUtf8StringDescriptor(libusb_device_handle *devh,
        uint8_t descIndex, uint8_t *data, int length);
    int GetStringDescriptorAscii(libusb_device_handle *devh,
        uint8_t descIndex, uint8_t *data, int length);
    int ProbeConfiguration(libusb_device *dev, libusb_device_descriptor &desc);
    void ProbeConfigurationSub1(const struct libusb_interface_descriptor *&intf, int &dfuMode,
        const struct UsbDfu::UsbDfuFuncDescriptor &funcDfu, libusb_device_descriptor const &desc,
        const struct libusb_config_descriptor *const &cfg) const;
    void ProbeConfigurationSub2(const libusb_device_descriptor &desc, const uint16_t &quirks,
        int &ret, libusb_device_handle *&devh, char (&serialName)[MAX_DESC_STR_LEN + 1]);
    int ProbeDevices(libusb_context *mCtx);
    void DisconnectDevices(void);
    int DfuseMultipleAlt(DfuIf &dfuseRoot) const;
    int DfuseDoBinDnload(struct DfuIf *dif, int xferSize,
                         const DfuFile &mFile, unsigned int startAddress);
    struct memsegment *FindSegment(struct memsegment *segmentList, unsigned int address) const;
    int AddSegment(struct memsegment **segmentList, struct memsegment segment) const;
    uint16_t GetQuirks(uint16_t vendor, uint16_t product, uint16_t bcdDevice) const;

    void HandleImageOTP(const BIN_INFO &info);
    void HandleImageVer(const BIN_INFO &info);
    void HandleImageFactory(const BIN_INFO &info, bool &isEraseAll);
    void HandleImageKVNV(const BIN_INFO &info, bool &isEraseAll);
    void HandleImageOthers(const BIN_INFO &info, bool &isEraseAll);
    void AddCmdByInfo(BIN_INFO info, bool &isEraseAll, QMap<int, BIN_INFO> &infoArray);
    void SetProgressRange() override;
    void UpdateProgress(UINT progress) override;
    void UpdateProgress(U32 fileSize, UINT totalFileSize = 0);
    void HidEnumerate();
    HANDLE OpenDevice(const CHAR *path, BOOL openRw) const;
    QString FindUsbDevice() const;
    int PreparationForBurning();
    int PartOfPreparationForBurning();
    void LoaderSendPayload(QString &strSendData, QMap<int, BIN_INFO> &tempArray, const BIN_INFO &info, int index);
    CHAR* HidEnumerateGetInterface(GUID interfaceClassGuid);
    bool IsUpgradeCompleted();
    void ConnectHidDfu();
    void HidDfuWriteToPortPartOne(int &retVal, int &overlapped, bool isLastPacket,
        unsigned short crc, OVERLAPPED &readOl) const;
    int HidDfuWriteToPort(const uint8_t *data, int len, bool isLastPacket, unsigned short crc) const;
    int DfuloadDoDnloadHidDfu(struct DfuFile &mFile);
    QString GetUSBPortChain(libusb_device* const dev) const;
    QString GetAllUSBInstanceId(QString location) const;
    bool IsResetCmd(const unsigned char *const cmdBuf, int len) const;
    QString SerialNameToDieID(const QString& serialName) const;
private:
    int dfuTimeout = 5000;              /* 5 seconds - default */
    unsigned int lastErasedPage = 1;    /* non-aligned value, won't match */
    unsigned int mDfuseAddress = 0;
    unsigned int mDfuseAddressPresent = 0;
    unsigned int mDfuseLength = 0;
    int mDfuseForce = 0;
    int mDfuseLeave = 0;
    int mDfuseUnprotect = 0;
    int mDfuseMassErase = 0;
    int mDfuseWillReset = 0;
    int verbose = 0;
    struct DfuIf *dfuRoot = nullptr;
    int matchVendor = -1;
    int matchProduct = -1;

    unsigned int mTransferSize = TRANSFER_SIZE;
    DfuMode m_mode = DfuMode::MODE_NONE;
    struct DfuStatus dfuStatus;
    libusb_context *ctx = nullptr;
    struct DfuFile file;
    int retValue;
    int dfuseDevice = 0;
    const char *mDfuseOptions = nullptr;
    uint16_t runtimeVendor;
    uint16_t runtimeProduct;
    unsigned short transaction = 0;
    bool isWriteLoader = false;
    HANDLE hidDeviceHandle = nullptr;
    int burnCmdReCount = 0;
    QString m_serialName;
    QString m_configPort;
};
}

#endif // BURNCTRLDFU_H
